"""
Created on 2021/12/15
note  : Predict and test with VGG19 model
author: Yuze Xuan, Xiaohu Hao, Xuan Wang, Sida Wang
"""

# * 导入包
import glob
import re
import shutil

import cv2
import numpy as np
import sklearn.metrics as metrics
from tensorflow.keras.models import load_model

from config import *


# * 图像加载函数
def load_image(image_path: str) -> np.ndarray:
    """ 图像加载.
    @param image_path: 图像路径.
    @return: 图像数据.
    """
    img_data = cv2.imread(image_path)
    img_data = cv2.resize(img_data, dsize=IMAGE_SIZE, interpolation=cv2.INTER_AREA)
    img_data = img_data.astype("float32")
    img_data /= 255.
    return np.array(img_data[:, :, :CHANNELS])


# * 加载模型权重
# 选择最优模型权重：多级排序-按val_acc升序，val_loss降序排序
weight_list = glob.glob(f'{SAVE_BASE_PATH}/checkpoints/*.hdf5')
weight_list = sorted(weight_list, key=lambda x: (-float(re.search(r'(.*)-(.*)-(.*).hdf5', x).group(3)),
                                                 float(re.search(r'(.*)-(.*)-(.*).hdf5', x).group(2))))
model = load_model(weight_list[0])
model.summary()

# * 加载测试数据集信息
LABELS = ['auditorium', 'football', 'gym', 'swimming_pool']
PREDICT_RESULT_BASE_PATH = os.path.join(SAVE_BASE_PATH, 'predict')
if SAVE_PRED_RESULT_BY_CAT:
    for label in LABELS:  # 为预测结果分类别建立存储路径
        if not os.path.exists(os.path.join(PREDICT_RESULT_BASE_PATH, label)):
            os.makedirs(os.path.join(PREDICT_RESULT_BASE_PATH, label))
true_labels, pred_labels = [], []
with open(os.path.join(DATASET_BASE_PATH, TEST_INFO_NAME), 'r+') as f:
    lines = f.readlines()

# * 图像加载及预测
for line in lines:
    if len(line) < 10:  # 排除空行
        continue
    else:
        file_name, true_label = line.strip().split(',')
        img_path = os.path.join(os.path.join(DATASET_BASE_PATH, 'test'), file_name)
        img = load_image(img_path)
        true_labels.append(true_label)
        pred_label = LABELS[np.argmax(model.predict(np.array([img])))]
        pred_labels.append(pred_label)
        if SAVE_PRED_RESULT_BY_CAT:
            # 将图像按预测结果保存到对应文件夹中
            shutil.copy(img_path, os.path.join(PREDICT_RESULT_BASE_PATH, pred_label))

# * 预测精度计算
acc = round(metrics.precision_score(true_labels, pred_labels, average='macro'), 4)
recall = round(metrics.recall_score(true_labels, pred_labels, average='macro'), 4)
f1_score = round(metrics.f1_score(true_labels, pred_labels, average='macro'), 4)

print("Test acc:{}, Test recall:{}, Test F1_score:{}".format(acc, recall, f1_score))
